#include "model/qwen2_hub.h"
#include "vocab/vocab.h"

#include <gtest/gtest.h>
#include <string>
#include <iostream>

using namespace pa;

TEST(VocabTest, TestVocab) {
    std::string hub = "data/Qwen2-0.5B";
    Qwen2Hub hub_model(ModelBase::MODEL_QWEN2_HUB, hub);
    bool flag = hub_model.is_valid();
    EXPECT_TRUE(flag);
    flag = hub_model.load();
    EXPECT_TRUE(flag);

    Vocab vocab(hub_model);
    flag = vocab.encoder_init();
    EXPECT_TRUE(flag);

    EXPECT_EQ(vocab.m_encoder.size(), 151643);
    EXPECT_EQ(vocab.m_added_token_encoder.size(), 3);
}

TEST(VocabTest, lut_test) {
    std::string hub = "data/Qwen2-0.5B";
    Qwen2Hub hub_model(ModelBase::MODEL_QWEN2_HUB, hub);
    bool flag = hub_model.is_valid();
    EXPECT_TRUE(flag);
    flag = hub_model.load();
    EXPECT_TRUE(flag);

    Vocab vocab(hub_model);
    flag = vocab.encoder_init();
    EXPECT_TRUE(flag);

    EXPECT_EQ(vocab.m_encoder.size(), 151643);
    EXPECT_EQ(vocab.m_added_token_encoder.size(), 3);

    const Vocab::Token& token = vocab.lookup(151644);
    EXPECT_EQ(token, "<|im_start|>");

    const auto& token2 = vocab.lookup(151658);
    EXPECT_EQ(token2, "");

    Vocab::Token token3 = "Hello";
    Vocab::Id idx = vocab.lookup(token3);
    EXPECT_EQ(idx, 19615);

    token3 = "中文的";
    idx = vocab.lookup(token3);
    EXPECT_EQ(idx, -1);

    token3 = "中";
    idx = vocab.lookup(token3);
    EXPECT_EQ(idx, -1);
}